# -*- coding: utf-8 -*-
# @Time    : 2019/2/25 13:34
# @Author  : cj
# @File    : dataset.py
import os
from dataset.data_item import DataItem


class DataSet:
    """
    封装的数据成员有path：数据的文件夹，database：数据的来源，data_items：为生成器，每次产生数据库路径下的
    一对image和XML
    """

    def __init__(self, path, database='Unknown'):
        if path is str:
            if not os.path.exists(path):
                raise FileNotFoundError(path + " Not Found")
        self._path = path
        self._database = database
        self._data_items = self._load()

    def _load(self):
        """
        加载数据，根据path
        """
        path_list = []
        if type(self._path) is str:
            path_list = os.listdir(self._path)
            for idx, image_name in enumerate(path_list):
                if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
                    image_path = os.path.join(self._path, image_name)
                    xml_path = os.path.join(self._path, image_name.split('.')[0] + '.xml')
                    data_item = DataItem()
                    try:
                        data_item.load(image_path, xml_path=xml_path)
                        yield data_item
                    except FileNotFoundError:
                        print(image_path + " or " + xml_path + " Not Found")
                        continue
        else:
            for name in self._path:
                try:
                    path_list.append(name['image'])
                except FileNotFoundError:
                    raise FileNotFoundError(name['image'] + 'NotFound')
            for idx, image_name in enumerate(path_list):
                if image_name.split('.')[-1] == 'png' or image_name.split('.')[-1] == 'jpg':
                    image_path = image_name
                    xml_path = image_name.split('.')[0] + '.xml'
                    data_item = DataItem()
                    try:
                        data_item.load(image_path, xml_path=xml_path)
                        yield data_item
                    except FileNotFoundError:
                        print(image_path + " or " + xml_path + " Not Found")
                        continue
            # self._data_items.append(data_item)

    def save_to_dir(self, dir_path: str):
        """
        :param dir_path: 要存储的数据库的路径
        """
        for data_item in self._data_items:
            data_item.save(dir_path, self._database)

    def save_to_tfrecords(self, dir_path: str, label_map_path: str):
        pass

    def get_xmls(self):
        """
        :return:返回一个生成器，每次返回一个XML的信息
        """
        # xlms = []
        for dataitem in self._data_items:
            annotation = dataitem.annotation
            yield annotation

    def get_imgs(self):
        """
        :return:返回一个生成器，每次返回一个img的信息
        """
        # imgs = []
        for dataitem in self._data_items:
            img = dataitem.image
            yield img

    def get_boundingboxes(self):
        """
        :return:返回一个生成器，每次返回一个boundingbox的信息
        """
        # boundingboxes = []
        for dataitem in self._data_items:
            boundingboxe = dataitem.annotation.get_boundingboxes()
            yield boundingboxe

    def get_classnames(self):
        """
        :return:返回一个生成器，每次返回一个classname的信息
        """
        # classnames = []
        for dataitem in self._data_items:
            classname = dataitem.annotation.get_classnames()
            yield classname

    def get_database(self):
        """
        :return:返回数据来源
        """
        return self._database

    def get_data_items(self):
        """
        :return:返回一个生成器，每次返回一个img和，与img对应的XML的信息
        """
        return self._data_items

    def set_database(self, database: str):
        """
        设置数据来源
        :param database:更改后的数据来源
        """
        self._database = database
